package cmd

import (
	"context"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/fsnotify/fsnotify"
	"github.com/spf13/cobra"
	admissionv1 "k8s.io/api/admission/v1"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/serializer"
)

type ResourceJsonPatch struct {
	Op    string        `json:"op"`
	Path  string        `json:"path"`
	Value corev1.EnvVar `json:"value"`
}

var (
	certDir string
	tlsCert string
	tlsKey  string
	port    int
	codecs  = serializer.NewCodecFactory(runtime.NewScheme())
	logger  = log.New(os.Stdout, "http: ", log.LstdFlags)
)

var rootCmd = &cobra.Command{
	Use:   "mutating-webhook",
	Short: "Kubernetes mutating webhook example",
	Long: `Example showing how to implement a basic mutating webhook in Kubernetes.

Example:
$ mutating-webhook --cert-dir <cert_dir> --tls-cert <tls_cert> --tls-key <tls_key> --port <port>`,
	Run: func(cmd *cobra.Command, args []string) {
		if tlsCert == "" || tlsKey == "" {
			fmt.Println("--tls-cert and --tls-key required")
			os.Exit(1)
		}
		runWebhookServer(tlsCert, tlsKey)
	},
}

// Execute adds all child commands to the root command and sets flags appropriately.
// This is called by main.main(). It only needs to happen once to the rootCmd.
func Execute() {
	cobra.CheckErr(rootCmd.Execute())
}

func init() {
	rootCmd.Flags().StringVar(&certDir, "cert-dir", "/var/lib/k8s-webhook-setenvresources/certs", "Directory containing TLS certificate and key")
	rootCmd.Flags().StringVar(&tlsCert, "tls-cert", "", "Certificate for TLS")
	rootCmd.Flags().StringVar(&tlsKey, "tls-key", "", "Private key file for TLS")
	rootCmd.Flags().IntVar(&port, "port", 443, "Port to listen on for HTTPS traffic")
}

func admissionReviewFromRequest(r *http.Request, deserializer runtime.Decoder) (*admissionv1.AdmissionReview, error) {
	// Validate that the incoming content type is correct.
	if r.Header.Get("Content-Type") != "application/json" {
		return nil, fmt.Errorf("expected application/json content-type")
	}

	// Get the body data, which will be the AdmissionReview
	// content for the request.
	var body []byte
	if r.Body != nil {
		requestData, err := io.ReadAll(r.Body)
		if err != nil {
			return nil, err
		}
		body = requestData
	}

	// Decode the request body into
	admissionReviewRequest := &admissionv1.AdmissionReview{}
	if _, _, err := deserializer.Decode(body, nil, admissionReviewRequest); err != nil {
		return nil, err
	}

	return admissionReviewRequest, nil
}

func mutatePod(w http.ResponseWriter, r *http.Request) {
	wanted_keys := []string{"POD_CPU_LIMIT", "POD_CPU_REQUEST", "POD_MEM_LIMIT", "POD_MEM_REQUEST"}
	cpu_limit_patch := ResourceJsonPatch{
		Op:   "add",
		Path: "/spec/containers/0/env/-",
		Value: corev1.EnvVar{
			Name: "POD_CPU_LIMIT",
			ValueFrom: &corev1.EnvVarSource{
				ResourceFieldRef: &corev1.ResourceFieldSelector{
					Resource: "limits.cpu",
				},
			},
		},
	}
	cpu_request_patch := ResourceJsonPatch{
		Op:   "add",
		Path: "/spec/containers/0/env/-",
		Value: corev1.EnvVar{
			Name: "POD_CPU_REQUEST",
			ValueFrom: &corev1.EnvVarSource{
				ResourceFieldRef: &corev1.ResourceFieldSelector{
					Resource: "requests.cpu",
				},
			},
		},
	}
	mem_limit_patch := ResourceJsonPatch{
		Op:   "add",
		Path: "/spec/containers/0/env/-",
		Value: corev1.EnvVar{
			Name: "POD_MEM_LIMIT",
			ValueFrom: &corev1.EnvVarSource{
				ResourceFieldRef: &corev1.ResourceFieldSelector{
					Resource: "limits.memory",
				},
			},
		},
	}
	mem_request_patch := ResourceJsonPatch{
		Op:   "add",
		Path: "/spec/containers/0/env/-",
		Value: corev1.EnvVar{
			Name: "POD_MEM_REQUEST",
			ValueFrom: &corev1.EnvVarSource{
				ResourceFieldRef: &corev1.ResourceFieldSelector{
					Resource: "requests.memory",
				},
			},
		},
	}

	wanted := map[string]ResourceJsonPatch{
		"POD_CPU_LIMIT":   cpu_limit_patch,
		"POD_CPU_REQUEST": cpu_request_patch,
		"POD_MEM_LIMIT":   mem_limit_patch,
		"POD_MEM_REQUEST": mem_request_patch,
	}

	logger.Printf("received message on mutate")

	deserializer := codecs.UniversalDeserializer()

	// Parse the AdmissionReview from the http request.
	admissionReviewRequest, err := admissionReviewFromRequest(r, deserializer)
	if err != nil {
		msg := fmt.Sprintf("error getting admission review from request: %v", err)
		logger.Printf(msg)
		w.WriteHeader(400)
		w.Write([]byte(msg))
		return
	}

	// Do server-side validation that we are only dealing with a pod resource. This
	// should also be part of the MutatingWebhookConfiguration in the cluster, but
	// we should verify here before continuing.
	podResource := metav1.GroupVersionResource{Group: "", Version: "v1", Resource: "pods"}
	if admissionReviewRequest.Request.Resource != podResource {
		msg := fmt.Sprintf("did not receive pod, got %s", admissionReviewRequest.Request.Resource.Resource)
		logger.Printf(msg)
		w.WriteHeader(400)
		w.Write([]byte(msg))
		return
	}

	// Decode the pod from the AdmissionReview.
	rawRequest := admissionReviewRequest.Request.Object.Raw
	pod := corev1.Pod{}
	if _, _, err := deserializer.Decode(rawRequest, nil, &pod); err != nil {
		msg := fmt.Sprintf("error decoding raw pod: %v", err)
		logger.Printf(msg)
		w.WriteHeader(500)
		w.Write([]byte(msg))
		return
	}

	// Create a response that will add a label to the pod if it does
	// not already have a label with the key of "hello". In this case
	// it does not matter what the value is, as long as the key exists.
	admissionResponse := &admissionv1.AdmissionResponse{}
	var no_needed_keys []string
	patchType := admissionv1.PatchTypeJSONPatch

	// if pod.Spec.Containers[0].Env == nil {

	// }
	envs := pod.Spec.Containers[0].Env
	for i := 0; i < len(envs); i++ {
		if strings.Contains(strings.Join(wanted_keys, ","), envs[i].Name) {
			no_needed_keys = append(no_needed_keys, envs[i].Name)
		}
	}

	var needed_keys []string
	for i := 0; i < len(wanted_keys); i++ {
		if !strings.Contains(strings.Join(no_needed_keys, ","), wanted_keys[i]) {
			needed_keys = append(needed_keys, wanted_keys[i])
		}
	}

	var patch []ResourceJsonPatch
	for i := 0; i < len(needed_keys); i++ {
		patch = append(patch, wanted[needed_keys[i]])
	}
	patch_bytes, err := json.Marshal(patch)
	if err != nil {
		logger.Printf(err.Error())
	}

	admissionResponse.Allowed = true
	if string(patch_bytes) != "" {
		admissionResponse.PatchType = &patchType
		admissionResponse.Patch = patch_bytes
	}

	// Construct the response, which is just another AdmissionReview.
	var admissionReviewResponse admissionv1.AdmissionReview
	admissionReviewResponse.Response = admissionResponse
	admissionReviewResponse.SetGroupVersionKind(admissionReviewRequest.GroupVersionKind())
	admissionReviewResponse.Response.UID = admissionReviewRequest.Request.UID

	resp, err := json.Marshal(admissionReviewResponse)
	if err != nil {
		msg := fmt.Sprintf("error marshalling response json: %v", err)
		logger.Printf(msg)
		w.WriteHeader(500)
		w.Write([]byte(msg))
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.Write(resp)
}

func runWebhookServer(certFile, keyFile string) {
	// cert, err := tls.LoadX509KeyPair(certFile, keyFile)
	// if err != nil {
	// 	panic(err)
	// }
	log.Println("Version: 1.2.2")

	fmt.Println("Starting webhook server")
	mux := http.NewServeMux()
	mux.HandleFunc("/mutate", mutatePod)
	server := http.Server{
		Addr:    fmt.Sprintf(":%d", port),
		Handler: mux,
		TLSConfig: &tls.Config{
			GetCertificate: func(hello *tls.ClientHelloInfo) (*tls.Certificate, error) {
				// This function is called for each new TLS connection
				// to determine the certificate to use.
				// tlsConfigMu.Lock()
				// defer tlsConfigMu.Unlock()
				return loadCertificate(certFile, keyFile), nil
			},
			// Optionally configure other TLS parameters here
			// MinVersion: tls.VersionTLS12,
			// CipherSuites: []uint16{tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256},
		},
		// TLSConfig: &tls.Config{
		// 	Certificates: []tls.Certificate{cert},
		// },

		ErrorLog: logger,
	}
	// 启动HTTP服务器
	go func() {
		if err := server.ListenAndServeTLS("", ""); err != http.ErrServerClosed {
			log.Fatalf("ListenAndServe error: %v", err)
		}
	}()

	// if err := server.ListenAndServeTLS("", ""); err != nil {
	// 	panic(err)
	// }

	// Handle signals to gracefully stop watching
	interrupt := make(chan os.Signal, 1)
	signal.Notify(interrupt, os.Interrupt, syscall.SIGTERM)
	defer signal.Stop(interrupt)
	// <-interrupt
	// go waitForShutdown(&server)

	// Create a new FSNotify watcher
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		log.Fatal(err)
	}
	defer watcher.Close()

	// Watch the specified directory
	err = watcher.Add(certDir)
	if err != nil {
		log.Fatal(err)
	}

	// Start watching for events
	for {
		select {
		case event, ok := <-watcher.Events:
			if !ok {
				return
			}
			if event.Has(fsnotify.Write) {
				// Update the certificate when the file is modified
				log.Println("File changed: ", event.Name)
				updateCertificate(&server, certFile, keyFile)
			}
		case err, ok := <-watcher.Errors:
			if !ok {
				return
			}
			log.Println("Error:", err)
		case <-interrupt:
			waitForShutdown(&server)
			return
		}
	}
}

func loadCertificate(certFile, keyFile string) *tls.Certificate {
	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		log.Fatalf("Failed to load certificate: %v", err)
	}
	return &cert
}

// Function to update the certificate dynamically
func updateCertificate(server *http.Server, certFile, keyFile string) {
	var tlsConfigMu sync.Mutex
	newCert, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		log.Printf("Failed to load new certificate: %v", err)
		return
	}

	// Update the server's TLSConfig with the new certificate
	tlsConfigMu.Lock()
	server.TLSConfig.Certificates = []tls.Certificate{newCert}
	tlsConfigMu.Unlock()

	log.Println("Certificate updated successfully")
}

func waitForShutdown(server *http.Server) {

	log.Println("Receive signal SIGTERM, Shutting down...")

	// Create a deadline to wait for
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	// Shutdown the server with the given context
	if err := server.Shutdown(ctx); err != nil {
		log.Fatalf("Server Shutdown Failed:%+v", err)
	}

	log.Println("Server exiting")
}
